package com.seo.textgen.pos;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import com.alex.utils.common.Random;
import com.seo.textgen.pos.tagger.TaggedSentence;
import com.seo.textgen.pos.tagger.Token;

public class Lexicon {

	private Map<String, Token[]> data = new HashMap<String, Token[]>();

	public void create(Collection<TaggedSentence> sentences) {
		Map<String, Set<Token>> temp = new HashMap<String, Set<Token>>();
		for (TaggedSentence sentence : sentences) {
			for (Token token : sentence.getTokens()) {
				Set<Token> tokens = temp.get(token.getTag());
				if (tokens == null) {
					tokens = new HashSet<Token>();
					temp.put(token.getTag(), tokens);
				}
				tokens.add(token);
			}
		}

		for (Map.Entry<String, Set<Token>> entry : temp.entrySet()) {
			data.put(entry.getKey(), entry.getValue().toArray(new Token[entry.getValue().size()]));
		}
	}

	public Token getRandomSubstitution(Token token) {
		Token[] tokens = data.get(token.getTag());
		if (tokens != null && tokens.length != 0) {
			return tokens[Random.getInt(tokens.length)];
		} else {
			return token;
		}
	}

	@Override
	public String toString() {
		return data.toString();
	}
}
